package cn.kennylee.codehub.springboot.microtrace.filter;

import cn.kennylee.codehub.springboot.microtrace.TraceProperties;
import cn.kennylee.codehub.springboot.microtrace.TraceUtils;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.Ordered;
import org.springframework.core.annotation.Order;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.lang.NonNull;
import org.springframework.util.StringUtils;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import reactor.core.publisher.Mono;

import static cn.kennylee.codehub.springboot.microtrace.TraceContext.generateTraceId;


/**
 * <p> Reactive微服务请求跟踪过滤器 </p>
 * <p>Created on 2025/4/30.</p>
 *
 * @author kennylee
 * @since 0.0.1
 */
@Slf4j
@Order(Ordered.HIGHEST_PRECEDENCE)
@RequiredArgsConstructor
public class ReactiveMicroTraceFilter implements WebFilter, MicroTraceFilter {

    private final TraceProperties traceProperties;

    @Override
    @NonNull
    public Mono<Void> filter(@NonNull ServerWebExchange exchange, @NonNull WebFilterChain chain) {
        // 在响应式编程环境中，不能使用 MDC（Mapped Diagnostic Context），因为存在根本缺陷
        // 每个请求都在不同的线程中执行，MDC 的数据不会自动传递到新的线程中
        // 因此，ThreadLocal 的值在异步操作后会丢失或者导致线程池污染（线程复用时携带旧值），内存泄漏等
        return Mono.deferContextual(contextView -> {
            // 检查是否已经存在 Trace ID
            if (contextView.hasKey(traceProperties.getMdcTraceIdName())) {
                log.warn("已经存在 Trace ID: {}", (String) contextView.get(traceProperties.getMdcTraceIdName()));
                return chain.filter(exchange);
            }

            // 1. 获取或生成 Trace ID
            String traceId = resolveOrGenerateTraceId(exchange);

            // 2. 修改请求头，添加 Trace ID
            ServerHttpRequest mutatedRequest = exchange.getRequest().mutate()
                .header(traceProperties.getTraceHeaderName(), traceId)
                .build();

            // 3. 重建 exchange 并处理上下文
            return chain.filter(exchange.mutate().request(mutatedRequest).build())
                // 4. 将 Trace ID 放入上下文中
                .contextWrite(ctx -> ctx.put(traceProperties.getMdcTraceIdName(), traceId))
                .doOnEach(signal -> {
                    // 响应式安全的日志记录
                    if (signal.isOnNext() || signal.isOnComplete()) {
                        log.trace("请求完成:[TraceID={}]", traceId);
                    }
                });
        });

    }

    /**
     * 解析或生成 Trace ID
     * <ol>
     *     <li>从请求头获取 Trace ID</li>
     *     <li>如果没有 Trace ID，则生成一个新的</li>
     * </ol>
     *
     * @param exchange ServerWebExchange
     * @return Trace ID
     */
    private String resolveOrGenerateTraceId(@NonNull ServerWebExchange exchange) {
        // 优先从请求头获取
        String traceId = TraceUtils.getTraceId(exchange);

        // 处理是否接收传递的 Trace ID
        traceId = TraceUtils.acceptTransferTraceId(traceId);

        if (!StringUtils.hasLength(traceId)) {
            traceId = generateTraceId();
            log.debug("从请求头找不到 Trace ID，新生成: {}", traceId);
        }
        return traceId;
    }
}
